package com.cloud.sample.stomp;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.context.event.EventListener;
import org.springframework.messaging.MessageHeaders;
import org.springframework.messaging.handler.annotation.DestinationVariable;
import org.springframework.messaging.simp.annotation.SubscribeMapping;
import org.springframework.messaging.support.GenericMessage;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.messaging.SessionConnectedEvent;
import org.springframework.web.socket.messaging.SessionDisconnectEvent;

import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;

@Slf4j
@Component
public class StompSessionManager {
    private Map<String, String> imUserIdMap = new HashMap<>();
    private Map<String, String> sessionIdMap = new HashMap<>();

    /**
     * im账号是否存在
     *
     * @param imUserId
     * @return
     */
    public boolean exist(String imUserId) {
        return imUserIdMap.containsKey(imUserId);
    }

    /**
     * 订阅
     *
     * @param imUserId
     */
    @SubscribeMapping(StompProvider.SUBSCRIBE_PREFIX + "{imUserId}")
    public void subscribe(@DestinationVariable("imUserId") String imUserId) {
        log.info("imUserId={} subscribe", imUserId);
    }

    /**
     * 连接开始
     *
     * @param event
     */
    @EventListener(SessionConnectedEvent.class)
    public void handleConnect(SessionConnectedEvent event) {
        MessageHeaders headers = event.getMessage().getHeaders();
        String sessionId = headers.get("simpSessionId").toString();
        String imUserId = getImUserId(headers.get("simpConnectMessage", GenericMessage.class));
        sessionIdMap.put(sessionId, imUserId);
        imUserIdMap.put(imUserId, sessionId);
    }

    /**
     * 连接关闭
     *
     * @param event
     */
    @EventListener(SessionDisconnectEvent.class)
    public void handleDisconnect(SessionDisconnectEvent event) {
        String imUserId = sessionIdMap.remove(event.getSessionId());
        if (StringUtils.isNotEmpty(imUserId)) {
            imUserIdMap.remove(imUserId);
        }
    }

    /**
     * 获取im账号
     *
     * @param message
     * @return
     */
    private String getImUserId(GenericMessage<String> message) {
        Map nativeHeaders = message.getHeaders().get("nativeHeaders", Map.class);
        return ((LinkedList) nativeHeaders.get("imUserId")).getFirst().toString();
    }
}
